#include "CrnnNet.h"
#include "OcrUtils.h"
#include <numeric>
#include <fstream>
#include <iostream>

CrnnNet::~CrnnNet() {
    net.clear();
}

void CrnnNet::setNumThread(int numOfThread) {
    numThread = numOfThread;
}

std::string readKeysFromAssets(const char* modelPath) {
    FLogger("readKeysFromAssets start...");
    if (modelPath == NULL) {
        FLogger(" %s", "AAssetManager==NULL");
        return NULL;
    }
    /*获取文件名并打开*/
    std::string path = modelPath;
    std::ifstream file(path, std::ios::binary | std::ios::ate);
    if (file.is_open()) {
        // 获取文件长度
        std::streampos fileSize = file.tellg();
        file.seekg(0, std::ios::beg);
        
        std::string buffer(fileSize, '\0');
        if (file.read(&buffer[0], fileSize)) {
            std::cout << buffer << std::endl;
        } else {
            std::cerr << "Error reading file " << path << std::endl;
        }
        file.close();
        return buffer;
    } else {
        std::cerr << "Unable to open file " << path << std::endl;
    }
    return "";
}

bool CrnnNet::initModel(const char* parmPath, const char* binPath, const char* keyPath) {
    int ret_param = net.load_param(parmPath);
    int ret_bin = net.load_model(binPath);
    if (ret_param != 0 || ret_bin != 0) {
        FLogger("# %d  %d", ret_param, ret_bin);
        return false;
    }

    std::string buffer = readKeysFromAssets(keyPath);
    if (!buffer.empty()) {
        std::istringstream inStr(buffer);
        std::string line;
        int size = 0;
        while (getline(inStr, line)) {
            keys.emplace_back(line);
            size++;
        }
        FLogger("keys size(%d)", size);
    } else {
        FLogger(" txt file not found");
        return false;
    }

    return true;
}

TextLine CrnnNet::scoreToTextLine(const float *srcData, int h, int w) {
    unsigned long keySize = keys.size();
    std::string strRes;
    std::vector<float> scores;
    int lastIndex = 0;

    for (int i = 0; i < h; i++) {
        int maxIndex = 0;
        float maxValue = -1000.f;
        //do softmax
        std::vector<float> exps(w);
        for (int j = 0; j < w; j++) {
            float e = srcData[i * w + j];
            float expSingle = exp(e);
            exps.at(j) = expSingle;
        }
        float partition = accumulate(exps.begin(), exps.end(), 0.0);//row sum
        for (int j = 0; j < w; j++) {
            float softmax = exps[j] / partition;
            if (softmax > maxValue) {
                maxValue = softmax;
                maxIndex = j;
            }
        }

        if (maxIndex > 0 && maxIndex < keySize && (!(i > 0 && maxIndex == lastIndex))) {
            scores.emplace_back(maxValue);
            strRes.append(keys[maxIndex - 1]);
        }
        lastIndex = maxIndex;
    }
    return {strRes, scores};
}

TextLine CrnnNet::getTextLine(const cv::Mat &src) {
    float scale = (float) dstHeight / (float) src.rows;
    int dstWidth = int((float) src.cols * scale);

    cv::Mat srcResize;
    cv::resize(src, srcResize, cv::Size(dstWidth, dstHeight));
    int type = srcResize.channels()==4? ncnn::Mat::PIXEL_RGBA2RGB : ncnn::Mat::PIXEL_RGB;
    ncnn::Mat input = ncnn::Mat::from_pixels(
            srcResize.data, type,
            srcResize.cols, srcResize.rows);

    input.substract_mean_normalize(meanValues, normValues);

    ncnn::Extractor extractor = net.create_extractor();
    extractor.set_num_threads(numThread);
    extractor.input("input", input);

    ncnn::Mat out;
    extractor.extract("out", out);

    return scoreToTextLine((float *) out.data, out.h, out.w);
}

std::vector<TextLine> CrnnNet::getTextLines(std::vector<cv::Mat> &partImg) {
    unsigned long size = partImg.size();
    std::vector<TextLine> textLines(size);
    for (int i = 0; i < size; ++i) {
        FLogger("get Text Line(%d)", i);
        //getTextLine
        double startCrnnTime = getCurrentTime();
        TextLine textLine = getTextLine(partImg[i]);
        double endCrnnTime = getCurrentTime();
        textLine.time = endCrnnTime - startCrnnTime;
        textLines[i] = textLine;
    }
    return textLines;
}
